package ext;

import util.TreeNode;

public class NodeDistance {
    
    //求一个 二叉树 里的2个节点的距离， 二叉树 每个节点的值是唯一的。
    public int dist(TreeNode root, int a, int b) {
        return dfs(root, a, b).dis;
    }

    private Msg dfs(TreeNode root, int a, int b) {
        if (root == null) {
            return new Msg(false,false, 0);
        }
        Msg left = dfs(root.left, a, b), 
            right = dfs(root.right, a, b);
        if (left.hasA && left.hasB) {
            return left;
        } else if (right.hasA && right.hasB) {
            return right;
        } else if (left.hasA && right.hasB || right.hasA && left.hasB) {
            return new Msg(true, true, left.dis + right.dis + root.val);
        } 
        boolean hasA = left.hasA || right.hasA || root.val == a;
        boolean hasB = left.hasB || right.hasB || root.val == b;
        int dis = 0;
        if (hasA || hasB) {
            dis = left.dis + right.dis + root.val;
        }
        return new Msg(hasA, hasB, dis);
    }

    class Msg {
        boolean hasA;
        boolean hasB;
        int dis;

        Msg(boolean hasA, boolean hasB, int dis) {
            this.hasA = hasA;
            this.hasB = hasB;
            this.dis = dis;
        }
    }

}
